import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
from detection.layers import weight_init

class VGG16BoxPredictor(nn.Module):
    def __init__(self, cfg, in_channels):
        super().__init__()
        num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES
        pool_size = cfg.MODEL.ROI_BOX_HEAD.POOL_RESOLUTION

        self.classifier = nn.Sequential(
            nn.Linear(in_channels * pool_size ** 2, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
        )

        self.cls_score = nn.Linear(4096, num_classes + 1)
        self.bbox_pred = nn.Linear(4096, num_classes * 4)
        nn.init.normal_(self.cls_score.weight, std=0.01)
        nn.init.normal_(self.bbox_pred.weight, std=0.001)
        for l in [self.cls_score, self.bbox_pred]:
            nn.init.constant_(l.bias, 0)

    def forward(self, box_features, is_da=False):
        box_features = box_features.view(box_features.size(0), -1)
        box_features = self.classifier(box_features)
        if is_da:
            return None, None, box_features

        class_logits = self.cls_score(box_features)
        box_regression = self.bbox_pred(box_features)
        return class_logits, box_regression, box_features

class SW_VGG16BoxPredictor(nn.Module):
    def __init__(self, cfg, in_channels):
        super().__init__()
        num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES
        pool_size = cfg.MODEL.ROI_BOX_HEAD.POOL_RESOLUTION
        self.gc = cfg.ADV.GC
        self.lc = cfg.ADV.LC
        self.da_use_contex = cfg.ADV.da_use_contex

        self.classifier = nn.Sequential(
            nn.Linear(in_channels * pool_size ** 2, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
        )

        feat_d = 4096
        if self.lc:
            feat_d += 128
        if self.gc:
            feat_d += 128
        self.cls_score = nn.Linear(feat_d, num_classes + 1)
        self.bbox_pred = nn.Linear(feat_d, num_classes * 4)
        nn.init.normal_(self.cls_score.weight, std=0.01)
        nn.init.normal_(self.bbox_pred.weight, std=0.001)
        for l in [self.cls_score, self.bbox_pred]:
            nn.init.constant_(l.bias, 0)

    def forward(self, box_features, feat_pixel, feat_gc, is_da=False):
        box_features = box_features.view(box_features.size(0), -1)
        box_features = self.classifier(box_features)

        instance_pooled_feat = box_features
        if self.lc:
            feat_pixel = feat_pixel.view(1, -1).repeat(box_features.size(0), 1)
            box_features = torch.cat((feat_pixel, box_features), 1)
            if self.da_use_contex:
                instance_pooled_feat = torch.cat((feat_pixel, instance_pooled_feat), 1)
        if self.gc:
            feat = feat_gc.view(1, -1).repeat(box_features.size(0), 1)
            box_features = torch.cat((feat, box_features), 1)
            if self.da_use_contex:
                instance_pooled_feat = torch.cat((feat, instance_pooled_feat), 1)

        if is_da:
            return None, None, box_features, instance_pooled_feat

        class_logits = self.cls_score(box_features)
        box_regression = self.bbox_pred(box_features)
        return class_logits, box_regression, box_features, instance_pooled_feat

class Predictor_C5(nn.Module):
    def __init__(self, cfg, in_channels):
        super().__init__()
        num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES

        resnet = models.resnet.__dict__[cfg.MODEL.BACKBONE.NAME](pretrained=True)#, norm_layer=FrozenBatchNorm2d)
        self._shared_roi_transform = resnet.layer4
        del resnet

        in_channels = self._shared_roi_transform[-1].conv3.out_channels
        self.cls_score = nn.Linear(in_channels, num_classes + 1)
        self.bbox_pred = nn.Linear(in_channels, num_classes * 4)
        nn.init.normal_(self.cls_score.weight, std=0.01)
        nn.init.normal_(self.bbox_pred.weight, std=0.001)
        for l in [self.cls_score, self.bbox_pred]:
            nn.init.constant_(l.bias, 0)

    def forward(self, box_features, is_da=False):
        box_features = self._shared_roi_transform(box_features)
        box_features = torch.mean(box_features, dim=(2, 3))
        if is_da:
            return None, None, box_features

        class_logits = self.cls_score(box_features)
        box_regression = self.bbox_pred(box_features)
        return class_logits, box_regression, box_features

class ConvFCBBoxHead(nn.Module):
    def __init__(self, cfg, in_channels):
        super().__init__()
        num_classes     = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES
        pool_size       = cfg.MODEL.ROI_BOX_HEAD.POOL_RESOLUTION

        self.shared_fcs = nn.Sequential(
            nn.Linear(in_channels * pool_size * pool_size, 1024),
            nn.Linear(1024, 1024)
        )
        self.cls_score = nn.Linear(1024, num_classes + 1)
        self.bbox_pred = nn.Linear(1024, num_classes * 4)

        # init
        for layer in self.shared_fcs:
            weight_init.c2_xavier_fill(layer)
        
        nn.init.normal_(self.cls_score.weight, std=0.01)
        nn.init.normal_(self.bbox_pred.weight, std=0.001)
        for l in [self.cls_score, self.bbox_pred]:
            nn.init.constant_(l.bias, 0)
    
    def forward(self, x):
        x = x.flatten(1)
        for fc in self.shared_fcs:
            x = F.relu(fc(x))
        # separate branches
        x_cls = x
        x_reg = x
        # box_features = x

        class_logits = self.cls_score(x_cls)
        box_regression = self.bbox_pred(x_reg)
        return class_logits, box_regression

BOX_PREDICTORS = {
    'vgg16_predictor': VGG16BoxPredictor,
    'r101_predictor_c5': Predictor_C5,
    'predictor_C5': Predictor_C5,
    'predictor_ConvFCBBoxHead':ConvFCBBoxHead,
}